"""
DetectModel: 基于 Ultralytics YOLO 的检测与目标确认/测距/调度框架
- 设计思想：将检测（detect）、目标识别（identifyObject）、测距（measureDistance）、
  位置/速度计算与管理器分离，提供清晰接口以便适配具体相机/机器人硬件。
- 用户需提供 camera, imu, robot 等接口（见 __init__ 中说明）。
"""

import time
import threading
from collections import deque, Counter
from typing import List, Tuple, Callable, Optional, Dict

from ultralytics import YOLO
import numpy as np

class DetectModel:
    def __init__(
        self,
        yolo_path: str,
        camera=None,
        imu=None,
        robot=None,
        conf_threshold: float = 0.8,
        result_buffer_size: int = 10,
        stale_clear_seconds: float = 2.0,
        imgsz: int = 640,
        camera_intrinsics: Optional[Dict[str, float]] = None,
    ):
        """
        Args:
            yolo_path: YOLO 模型权重路径（.pt/.yaml/...）
            camera: 相机接口对象（需实现 get_frame() -> np.ndarray, get_align_depth() -> depth_map）
                    — get_frame() 返回当前 RGB 帧
                    — get_align_depth() 返回与 RGB 对齐的深度图（同 shape，高度为米或毫米，需约定）
            imu: IMU 接口对象（需实现 get_imu() -> dict，包括角速度/加速度/姿态等）
            robot: 机器人控制接口对象（需实现 send_velocity(vx, vy, yaw_rate) 等，或 DogTurn/DogController）
            conf_threshold: 计入结果列表的概率阈值（你指定的 0.8）
            result_buffer_size: 结果列表满多少条时做一次结算（你指定 10）
            stale_clear_seconds: 若超过该时间未收到新高置信结果则清空结果列表（你指定 2s）
            imgsz: YOLO 输入尺寸
            camera_intrinsics: 可选相机内参字典，例如 {'fx':..., 'fy':..., 'cx':..., 'cy':...}
        """
        self.model = YOLO(yolo_path)
        self.camera = camera
        self.imu = imu
        self.robot = robot

        self.conft = conf_threshold
        self.buf_size = result_buffer_size
        self.stale_seconds = stale_clear_seconds
        self.imgsz = imgsz

        self.names = getattr(self.model, "names", None) or {}
        # 内部缓冲：保存 (label, prob, timestamp)
        self._result_buf = deque(maxlen=self.buf_size)
        self._last_insert_ts = 0.0
        self._lock = threading.Lock()

        # 摄像机内参（用于深度->空间坐标转换），如果 None，computeRelativePosition 会做最简单的近似
        self.K = camera_intrinsics or {"fx": 1.0, "fy": 1.0, "cx": 0.0, "cy": 0.0}

    # -----------------------------
    # 核心：detect(frame)
    # -----------------------------
    def detect(self, frame: np.ndarray) -> List[Tuple[str, float]]:
        """
        对单帧进行 YOLO 推断，筛选出置信度高于 self.conft 的目标并将其加入内部结果队列。
        额外功能：
         - 如果队列满（达到 self.buf_size），则对队列做结算并返回当前视角下出现次数最多的两个目标 [(label, prob), ...]
         - 如果距离上次插入超过 self.stale_seconds 且没有新高置信结果，则会清空队列
        返回：
            list of (label, prob) —— 如果触发结算返回两个最多的目标，否则返回当前检测到的高置信项列表（可能为空）
        """
        if frame is None:
            return []

        results_high = []  # 本次推断中满足阈值的 (label,prob)

        try:
            # Ultralytics 支持直接传入 numpy array 作为 source
            results = self.model.predict(source=frame, imgsz=self.imgsz, conf=0.001, verbose=False)
            # results 可能为列表，每项对应一帧（这里只传入单帧）
            if len(results) == 0:
                return []
            r = results[0]
            # boxes 可能存在不同属性，根据 ultralytics 版本差别处理
            try:
                boxes = getattr(r, "boxes", None)
                if boxes is None:
                    # 兼容旧版本：r.boxes might be list-like; try r.boxes.data
                    boxes = r.boxes
                # boxes.xyxy, boxes.conf, boxes.cls etc.
                confs = getattr(boxes, "conf", None)
                cls_vals = getattr(boxes, "cls", None)
                # boxes might be a tensor of shape (N,6) -> xyxy conf cls
                if confs is None or cls_vals is None:
                    # fallback: try boxes.data (tensor Nx6)
                    data = getattr(boxes, "data", None)
                    if data is not None:
                        arr = data.cpu().numpy()
                        for row in arr:
                            # row format may be [x1,y1,x2,y2,conf,cls]
                            if row.shape[0] >= 6:
                                prob = float(row[4])
                                cls_idx = int(row[5])
                                label = str(self.names.get(cls_idx, str(cls_idx)))
                                if prob >= self.conft:
                                    results_high.append((label, prob))
                    else:
                        # 无法解析 boxes，跳过
                        pass
                else:
                    # 常见情况： boxes.conf 和 boxes.cls 都可用（tensor）
                    confs_np = np.array(confs.cpu()) if hasattr(confs, "cpu") else np.array(confs)
                    cls_np = np.array(cls_vals.cpu()) if hasattr(cls_vals, "cpu") else np.array(cls_vals)
                    for ci, p in zip(cls_np, confs_np):
                        prob = float(p)
                        cls_idx = int(ci)
                        label = str(self.names.get(cls_idx, str(cls_idx)))
                        if prob >= self.conft:
                            results_high.append((label, prob))
            except Exception:
                # 最后一招：从 r.pred 或 r.boxes.xyxy/cls/conf 里尝试读取
                try:
                    # r.boxes.xyxy, r.boxes.conf, r.boxes.cls
                    for i in range(len(r.boxes)):
                        b = r.boxes[i]
                        prob = float(b.conf) if hasattr(b, "conf") else float(b[4])
                        cls_idx = int(b.cls) if hasattr(b, "cls") else int(b[5])
                        if prob >= self.conft:
                            label = str(self.names.get(cls_idx, str(cls_idx)))
                            results_high.append((label, prob))
                except Exception:
                    # 放弃本帧解析（保证鲁棒）
                    results_high = []

        except Exception as e:
            # 推断出错不崩溃，记录并返回空
            print(f"[DetectModel] YOLO 推断异常: {e}")
            return []

        # 将高置信结果加入缓冲（带锁保护）
        now = time.time()
        inserted = False
        with self._lock:
            if results_high:
                for label, prob in results_high:
                    self._result_buf.append((label, prob, now))
                self._last_insert_ts = now
                inserted = True
            else:
                # 检查是否需要清空缓冲（超过 stale_seconds 没有新条目）
                if self._result_buf and (now - self._last_insert_ts) > self.stale_seconds:
                    self._result_buf.clear()

        # 如果缓冲已满（等于 buffer_size）则结算并返回两个出现次数最多目标
        with self._lock:
            if len(self._result_buf) >= self.buf_size:
                # 统计 label 出现次数并取 top2。对同 label 取最高 prob 作为该 label 的代表置信度
                labels = [item[0] for item in self._result_buf]
                counts = Counter(labels)
                top_two = counts.most_common(2)
                results_out = []
                for label, _ in top_two:
                    # 在缓冲中找该 label 的最大 prob
                    probs_for_label = [p for (l, p, ts) in self._result_buf if l == label]
                    best_prob = max(probs_for_label) if probs_for_label else 0.0
                    results_out.append((label, best_prob))
                # 清空缓冲，准备下一轮
                self._result_buf.clear()
                self._last_insert_ts = 0.0
                return results_out

        # 未触发结算，返回本次检测到的高置信项（可能为空）
        return results_high

    # -----------------------------
    # identifyObject
    # -----------------------------
    def identifyObject(
        self,
        target_label: str,
        frame_provider: Optional[Callable[[], np.ndarray]] = None,
        timeout: float = 2.0,
        poll_interval: float = 0.1,
    ) -> Tuple[bool, str]:
        """
        持续比对 detect 的输出以确认目标是否出现在当前视角中。
        两种使用模式：
         - 若只传入一个 frame（frame_provider=None），将对该帧做一次检测并返回一次性结果。
         - 推荐：传入 frame_provider（可调用函数）用于持续拉帧（例如 camera.get_frame）。
           函数将最多等待 timeout 秒，若 detect 返回含 target_label 的条目则返回 (True, target_label)，
           否则超时返回 (False, target_label)。

        Args:
            target_label: 要识别的目标名称（字符串，与 model.names 对齐）
            frame_provider: 一个无参可调用对象，返回最新帧（np.ndarray）
            timeout: 最长等待时间（秒）
            poll_interval: 每次循环的等待时间（秒）
        """
        start = time.time()
        # 如果没有 frame_provider，但传入的是单帧函数：允许传入 lambda: frame 常量
        if frame_provider is None:
            # 只做一次 detect（适用于同步场景）
            return_any = False
            # 没有帧提供器无法继续，返回 False
            return False, target_label

        # 如果有 frame_provider，循环拉帧并检测
        while time.time() - start <= timeout:
            try:
                frame = frame_provider()
            except Exception as e:
                # 拉帧异常：睡一会儿再试
                print(f"[identifyObject] frame_provider 取帧异常: {e}")
                time.sleep(poll_interval)
                continue

            dets = self.detect(frame)
            # dets 是一个 list of (label, prob)
            for label, prob in dets:
                if label == target_label:
                    return True, target_label
            time.sleep(poll_interval)

        return False, target_label

    # -----------------------------
    # ManagerIdentify (调度函数)
    # -----------------------------
    def ManagerIdentify(self, target_label: str, frame_provider: Callable[[], np.ndarray]):
        """
        使用 identifyObject 结果做调度：
         - 如果识别成功 (True, object)，调用 measureDistance 进行测距；
         - 如果识别失败 (False, object)，调用 DogController('turn') 进行转向寻目标。
        该函数是同步的高层接口，调用者可在循环中反复调用或放到线程中运行。
        """
        ok, obj = self.identifyObject(target_label, frame_provider=frame_provider)
        if ok:
            # 获取当前帧并测距
            try:
                frame = frame_provider()
                res = self.measureDistance(obj, frame)
                return res  # (object, prob, distance) 或 None
            except Exception as e:
                print(f"[ManagerIdentify] measureDistance 异常: {e}")
                return None
        else:
            # 通知底层转向（字符串 'turn'）
            self.DogController('turn')
            return None

    # -----------------------------
    # DogTurn / DogController（底层留空或调用 robot 接口）
    # -----------------------------
    def DogTurn(self, angle: float):
        """
        保留给你实现的底层控制函数（你说不改这里），不过我们提供一个可选的 robot 接口适配：
        如果构造时传入了 self.robot 并带有 turn(angle) 方法，则调用它。
        """
        if hasattr(self, "robot") and self.robot is not None:
            try:
                if hasattr(self.robot, "turn"):
                    self.robot.turn(angle)
                elif hasattr(self.robot, "DogTurn"):
                    self.robot.DogTurn(angle)
                else:
                    # 尝试通过统一命令发送
                    if hasattr(self.robot, "send_command"):
                        self.robot.send_command({"cmd": "turn", "angle": angle})
            except Exception as e:
                print(f"[DogTurn] 调用 robot.turn 异常: {e}")

    def DogController(self, cmd: str):
        """
        统一的命令接口：若 cmd == 'turn' 则调用 DogTurn(180)
        你可替换或在外部实现 robot.send_command 以实现更复杂策略
        """
        if cmd == 'turn':
            # 这里使用字符串 '180' 在你原始代码里，但更合理用数字
            try:
                self.DogTurn(180)
            except Exception:
                # 不可用时尝试 robot 的 turn 方法
                if self.robot and hasattr(self.robot, "turn"):
                    self.robot.turn(180)

    # -----------------------------
    # measureDistance
    # -----------------------------
    def measureDistance(self, target_label: str, frame: np.ndarray, depth_threshold: float = 0.5):
        """
        - 对传入帧再次进行 YOLO 推断，筛选出 target_label 的 bbox 和 prob。
        - 如果 prob > 0.5（或阈值），则读取相机的 aligned depth 图（camera.get_align_depth）
          并计算 bbox 区域的平均深度作为距离返回。
        - 如果 prob < 0.5，则调用 identifyObject 再次确认（同步等待），若确认成功继续定位测距。
        返回：
            (object_label, prob, distance_in_meters) 或 None（失败）
        """
        try:
            # 1) 再次推断
            results = self.model.predict(source=frame, imgsz=self.imgsz, conf=0.001, verbose=False)
            if len(results) == 0:
                return None
            r = results[0]
            # 尝试遍历 boxes
            found = None
            try:
                boxes = getattr(r, "boxes", None)
                if boxes is not None:
                    # boxes tensor-like
                    confs = getattr(boxes, "conf", None)
                    cls_vals = getattr(boxes, "cls", None)
                    xyxy = getattr(boxes, "xyxy", None)
                    # iterate
                    if confs is not None and cls_vals is not None and xyxy is not None:
                        confs_np = np.array(confs.cpu()) if hasattr(confs, "cpu") else np.array(confs)
                        cls_np = np.array(cls_vals.cpu()) if hasattr(cls_vals, "cpu") else np.array(cls_vals)
                        xyxy_np = np.array(xyxy.cpu()) if hasattr(xyxy, "cpu") else np.array(xyxy)
                        for cls_idx, p, box_xy in zip(cls_np, confs_np, xyxy_np):
                            label = str(self.names.get(int(cls_idx), str(int(cls_idx))))
                            if label == target_label:
                                found = (label, float(p), box_xy)  # box_xy = [x1,y1,x2,y2]
                                break
            except Exception:
                # Fallback parsing
                try:
                    data = getattr(r.boxes, "data", None)
                    if data is not None:
                        arr = data.cpu().numpy()
                        for row in arr:
                            prob = float(row[4])
                            cls_idx = int(row[5])
                            label = str(self.names.get(cls_idx, str(cls_idx)))
                            if label == target_label:
                                box_xy = row[:4]
                                found = (label, prob, box_xy)
                                break
                except Exception:
                    found = None

            if not found:
                # 若未找到目标或找到但 prob < 0.5，则进一步调用 identifyObject 确认（同步，最多等待 self.stale_seconds）
                ok, _ = self.identifyObject(target_label, frame_provider=(lambda: frame), timeout=self.stale_seconds)
                if not ok:
                    return None
                # 再测一次（递归一次）
                return self.measureDistance(target_label, frame, depth_threshold)

            label, prob, box_xy = found
            if prob < 0.5:
                # 同上逻辑，确认后再测
                ok, _ = self.identifyObject(target_label, frame_provider=(lambda: frame), timeout=self.stale_seconds)
                if not ok:
                    return None
                return self.measureDistance(target_label, frame, depth_threshold)

            # 若 prob >= 0.5，读取深度图并求 bbox 区域平均深度
            if not self.camera or not hasattr(self.camera, "get_align_depth"):
                # 没有深度访问接口，返回 None 或者 NaN 距离
                return (label, prob, float('nan'))

            depth_map = self.camera.get_align_depth()  # 假定返回与 frame 同样 HxW 的数组（单位：米）
            if depth_map is None:
                return (label, prob, float('nan'))

            x1, y1, x2, y2 = map(int, box_xy[:4])
            # 限定到图像范围，避免越界
            h, w = depth_map.shape[:2]
            x1 = max(0, min(w - 1, x1))
            x2 = max(0, min(w - 1, x2))
            y1 = max(0, min(h - 1, y1))
            y2 = max(0, min(h - 1, y2))
            if x2 <= x1 or y2 <= y1:
                return (label, prob, float('nan'))

            roi = depth_map[y1:y2, x1:x2]
            # 过滤掉无效深度（例如 0 或 NaN）
            roi_valid = roi[np.isfinite(roi) & (roi > 0)]
            if roi_valid.size == 0:
                # 没有有效深度，尝试膨胀 bbox 周围采样（简单策略）
                pad = 5
                x1p, x2p = max(0, x1 - pad), min(w - 1, x2 + pad)
                y1p, y2p = max(0, y1 - pad), min(h - 1, y2 + pad)
                roi2 = depth_map[y1p:y2p, x1p:x2p]
                roi_valid = roi2[np.isfinite(roi2) & (roi2 > 0)]
                if roi_valid.size == 0:
                    return (label, prob, float('nan'))

            avg_depth = float(np.mean(roi_valid))
            return (label, prob, avg_depth)

        except Exception as e:
            print(f"[measureDistance] 异常: {e}")
            return None

    # -----------------------------
    # computeRelativePosition
    # -----------------------------
    def computeRelativePosition(self, obb, average_distance: float):
        """
        根据目标在图像中的像素坐标（obb，假设为 [x1,y1,x2,y2] 或中心点 [cx,cy]）和平均距离，
        使用针孔相机模型估计目标在相机坐标系中的相对位置 (x, y, z)。

        返回 (x, y, z)（单位与 depth 一致）
        """
        try:
            fx = self.K.get("fx", 1.0)
            fy = self.K.get("fy", 1.0)
            cx = self.K.get("cx", 0.0)
            cy = self.K.get("cy", 0.0)

            if len(obb) >= 4:
                # 以 bbox 中心作为像素坐标
                x1, y1, x2, y2 = obb[:4]
                cx_px = (x1 + x2) / 2.0
                cy_px = (y1 + y2) / 2.0
            else:
                # 直接传入中心
                cx_px, cy_px = obb[:2]

            Z = average_distance
            if not np.isfinite(Z) or Z <= 0:
                # 无效距离，返回 None
                return None

            X = (cx_px - cx) * Z / fx
            Y = (cy_px - cy) * Z / fy
            return (X, Y, Z)
        except Exception as e:
            print(f"[computeRelativePosition] 异常: {e}")
            return None

    # -----------------------------
    # computeVelocity (PID 框架)
    # -----------------------------
    def computeVelocity(self, target_pos: Tuple[float, float, float], pid_params=None, dt: float = 0.05):
        """
        根据相对位置 (X, Y, Z) 与 IMU 信息，使用简单 PID（这里只实现 P 控制 + 基础衰减）来计算机器狗的速度命令。
        假定机器狗支持 vx, vy, yaw_rate（全向底盘模型）。

        Args:
            target_pos: (X, Y, Z) 目标在相机/机器人坐标系中的相对位置（X:右向, Y:下向, Z:前向）
            pid_params: 可传入 {'Kp_x':..., 'Kp_y':..., 'Kp_yaw':...}
            dt: 控制周期估计

        Returns:
            (vel_x, vel_y, yaw_z)
        """
        if target_pos is None:
            return (0.0, 0.0, 0.0)

        params = pid_params or {}
        Kp_x = params.get("Kp_x", 0.8)
        Kp_y = params.get("Kp_y", 0.8)
        Kp_yaw = params.get("Kp_yaw", 1.0)

        X, Y, Z = target_pos
        # 目标角度（偏航）—— 简单使用横向误差与前向距离估计期望 yaw
        desired_yaw = np.arctan2(X, Z)  # 以弧度表示，右偏为正
        # 速度命令（P 控制）
        vx = Kp_x * Z  # 向前速度与深度正相关（可能需要饱和）
        vy = -Kp_y * X  # 想把机器人移到目标 X=0（注意坐标约定）
        yaw_rate = Kp_yaw * desired_yaw

        # 饱和与防爆控制（根据你的底盘能力调整）
        max_v = 1.0  # m/s
        max_yaw = 1.0  # rad/s
        vx = max(-max_v, min(max_v, vx))
        vy = max(-max_v, min(max_v, vy))
        yaw_rate = max(-max_yaw, min(max_yaw, yaw_rate))

        return (vx, vy, yaw_rate)

    # -----------------------------
    # GlobalManager
    # -----------------------------
    def GlobalManager(self, target_label: str, frame_provider: Callable[[], np.ndarray], distance_threshold: float = 0.1):
        """
        高层流程：
         1. 检查目标是否在当前视角（identifyObject）
         2. 如果在：测距 -> 计算相对位置 -> 计算速度 -> 下发速度命令；循环直到距离小于 distance_threshold
         3. 如果不在：调用 DogController('turn') 旋转机器人并继续寻找
        该函数是阻塞式的高层管理器，可放到线程中执行。
        """
        # 循环直到找到并接近目标
        while True:
            ok, _ = self.identifyObject(target_label, frame_provider=frame_provider, timeout=2.0)
            if not ok:
                # 旋转寻目标
                self.DogController('turn')
                # 给一点时间让机器人转动后再拉帧检查
                time.sleep(0.5)
                continue

            # 找到后测距与定位
            frame = frame_provider()
            res = self.measureDistance(target_label, frame)
            if res is None:
                # 尝试继续下一次识别
                time.sleep(0.1)
                continue

            label, prob, dist = res
            if not np.isfinite(dist):
                # 没有有效深度，尝试靠近一点（转到目标）或重试
                time.sleep(0.1)
                continue

            # 获取 bbox（再次推断以获得 bbox）
            # 这里简单调用 model.predict 并解析第一个匹配 bbox（同 measureDistance）
            results = self.model.predict(source=frame, imgsz=self.imgsz, conf=0.001, verbose=False)
            obb = None
            if results and len(results) > 0:
                r = results[0]
                try:
                    boxes = getattr(r, "boxes", None)
                    if boxes is not None:
                        confs = getattr(boxes, "conf", None)
                        cls_vals = getattr(boxes, "cls", None)
                        xyxy = getattr(boxes, "xyxy", None)
                        if confs is not None and cls_vals is not None and xyxy is not None:
                            confs_np = np.array(confs.cpu()) if hasattr(confs, "cpu") else np.array(confs)
                            cls_np = np.array(cls_vals.cpu()) if hasattr(cls_vals, "cpu") else np.array(cls_vals)
                            xyxy_np = np.array(xyxy.cpu()) if hasattr(xyxy, "cpu") else np.array(xyxy)
                            for cls_idx, p, box_xy in zip(cls_np, confs_np, xyxy_np):
                                label_now = str(self.names.get(int(cls_idx), str(int(cls_idx))))
                                if label_now == target_label:
                                    obb = box_xy
                                    break
                except Exception:
                    obb = None

            if obb is None:
                # 没拿到 bbox，则仅使用深度做一个简单向前移动
                pos = (0.0, 0.0, dist)
            else:
                pos = self.computeRelativePosition(obb, dist)

            # 计算速度并下发
            vel_cmd = self.computeVelocity(pos)
            if self.robot and hasattr(self.robot, "send_velocity"):
                try:
                    self.robot.send_velocity(*vel_cmd)
                except Exception as e:
                    print(f"[GlobalManager] 发送速度异常: {e}")
            else:
                # 如果没有 robot 接口，只打印或通过 DogController 适配
                print(f"[GlobalManager] 速度命令: vx={vel_cmd[0]:.3f}, vy={vel_cmd[1]:.3f}, yaw={vel_cmd[2]:.3f}")

            # 距离判断
            if dist <= distance_threshold:
                # 到达附近，停止并返回
                if self.robot and hasattr(self.robot, "send_velocity"):
                    try:
                        self.robot.send_velocity(0.0, 0.0, 0.0)
                    except Exception:
                        pass
                return True

            # 小的 sleep 限制循环频率
            time.sleep(0.05)


# -----------------------------
# 使用示例（伪代码）
# -----------------------------
if __name__ == "__main__":
    # 假设你自己实现了下面这些类并传入
    class DummyCamera:
        def get_frame(self):
            # 返回 np.ndarray
            return np.zeros((480, 640, 3), dtype=np.uint8)
        def get_align_depth(self):
            # 返回深度图（单位：米）
            return np.ones((480, 640), dtype=float) * 2.0

    cam = DummyCamera()
    model = DetectModel("yolov8n.pt", camera=cam)

    # frame_provider 示例
    def frame_provider():
        return cam.get_frame()

    # 识别示例（阻塞等待最多 2s）
    ok, obj = model.identifyObject("ball", frame_provider=frame_provider, timeout=2.0)
    print("识别结果：", ok, obj)

    # 全局管理器（会在找到目标并接近 0.1m 后返回）
    # model.GlobalManager("ball", frame_provider)
